# 服务器主程序
from mcp.server.fastmcp import FastMCP
import sys
import logging

logger = logging.getLogger('MCPServer')

# 修复Windows控制台的UTF-8编码问题
if sys.platform == 'win32':
    sys.stderr.reconfigure(encoding='utf-8')
    sys.stdout.reconfigure(encoding='utf-8')

import math
import random
import os
import shutil
import pygame
from pathlib import Path
import subprocess
from pywinauto import Application, Desktop
import psutil
import win32gui
import win32con
import tempfile
import json
import time
import ctypes
import win32api
import win32com.client
from time import sleep

# 初始化pygame混音器用于音乐播放
pygame.mixer.init()

# 创建MCP服务器
mcp = FastMCP("MCPServer", version="1.0.0")


# 添加计算工具
@mcp.tool()
def calculator(python_expression: str) -> dict:
    """用于数学计算，始终使用此工具计算Python表达式的结果。可以使用`math`和`random`模块。"""
    result = eval(python_expression)
    logger.info(f"Calculating formula: {python_expression}, result: {result}")
    return {"success": True, "result": result}

# 添加文件系统工具
@mcp.tool()
def create_directory(path: str) -> dict:
    """在指定路径创建新目录。"""
    try:
        os.makedirs(path, exist_ok=True)
        logger.info(f"Created directory: {path}")
        return {"success": True, "message": f"目录已创建: {path}"}
    except Exception as e:
        logger.error(f"创建目录失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def delete_directory(path: str) -> dict:
    """删除指定目录及其所有内容。"""
    try:
        shutil.rmtree(path)
        logger.info(f"Deleted directory: {path}")
        return {"success": True, "message": f"目录已删除: {path}"}
    except Exception as e:
        logger.error(f"删除目录失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def list_directory(path: str) -> dict:
    """列出指定路径中的所有文件和目录。"""
    try:
        items = os.listdir(path)
        logger.info(f"Listed directory contents: {path}")
        return {"success": True, "items": items}
    except Exception as e:
        logger.error(f"列出目录失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def create_file(path: str, content: str = "") -> dict:
    """创建新文件并可选择写入内容。"""
    try:
        with open(path, 'w', encoding='utf-8') as f:
            f.write(content)
        logger.info(f"Created file: {path}")
        return {"success": True, "message": f"文件已创建: {path}"}
    except Exception as e:
        logger.error(f"创建文件失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def read_file(path: str) -> dict:
    """读取文件内容。"""
    try:
        with open(path, 'r', encoding='utf-8') as f:
            content = f.read()
        logger.info(f"Read file: {path}")
        return {"success": True, "content": content}
    except Exception as e:
        logger.error(f"读取文件失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def write_file(path: str, content: str) -> dict:
    """向文件写入内容。"""
    try:
        with open(path, 'w', encoding='utf-8') as f:
            f.write(content)
        logger.info(f"Wrote to file: {path}")
        return {"success": True, "message": f"内容已写入: {path}"}
    except Exception as e:
        logger.error(f"写入文件失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def delete_file(path: str) -> dict:
    """删除文件。"""
    try:
        os.remove(path)
        logger.info(f"Deleted file: {path}")
        return {"success": True, "message": f"文件已删除: {path}"}
    except Exception as e:
        logger.error(f"删除文件失败: {str(e)}")
        return {"success": False, "error": str(e)}

# 添加音乐播放工具
@mcp.tool()
def play_music(file_path: str) -> dict:
    """播放音乐文件。"""
    try:
        if not os.path.exists(file_path):
            return {"success": False, "error": "文件不存在"}
            
        if not os.path.splitext(file_path)[1].lower() in ['.mp3', '.wav']:
            return {"success": False, "error": "不支持的文件格式"}
            
        if Path(file_path).stat().st_size == 0:
            return {"success": False, "error": "文件为空"}
            
        pygame.mixer.music.load(file_path)
        pygame.mixer.music.play()
        logger.info(f"正在播放: {os.path.basename(file_path)}")
        return {
            "success": True, 
            "message": f"正在播放: {os.path.basename(file_path)}",
            "file_info": {
                "name": os.path.basename(file_path),
                "path": file_path,
                "format": os.path.splitext(file_path)[1][1:]
            }
        }
    except Exception as e:
        logger.error(f"播放失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def pause_music() -> dict:
    """暂停当前播放的音乐。"""
    try:
        pygame.mixer.music.pause()
        logger.info("音乐已暂停")
        return {"success": True, "message": "音乐已暂停"}
    except Exception as e:
        logger.error(f"暂停音乐失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def unpause_music() -> dict:
    """恢复播放暂停的音乐。"""
    try:
        pygame.mixer.music.unpause()
        logger.info("音乐已恢复")
        return {"success": True, "message": "音乐已恢复"}
    except Exception as e:
        logger.error(f"恢复音乐失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def stop_music() -> dict:
    """停止当前播放的音乐。"""
    try:
        pygame.mixer.music.stop()
        logger.info("音乐已停止")
        return {"success": True, "message": "音乐已停止"}
    except Exception as e:
        logger.error(f"停止音乐失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def set_volume(volume: float) -> dict:
    """设置音乐音量（0.0到1.0）。"""
    try:
        if 0.0 <= volume <= 1.0:
            pygame.mixer.music.set_volume(volume)
            logger.info(f"音量设置为: {volume}")
            return {"success": True, "message": f"音量设置为: {volume}"}
        else:
            return {"success": False, "error": "音量必须在0.0到1.0之间"}
    except Exception as e:
        logger.error(f"设置音量失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def get_playback_status() -> dict:
    """获取当前播放状态。"""
    try:
        is_playing = pygame.mixer.music.get_busy()
        return {
            "success": True,
            "is_playing": is_playing,
            "volume": pygame.mixer.music.get_volume()
        }
    except Exception as e:
        logger.error(f"获取播放状态失败: {str(e)}")
        return {"success": False, "error": str(e)}

import os.path

# 添加全局变量来管理播放列表
current_playlist = []
current_track_index = -1

@mcp.tool()
def scan_music_directory(directory: str = "music") -> dict:
    """扫描音乐目录并返回所有MP3文件"""
    try:
        music_files = []
        base_path = os.path.join(os.path.dirname(__file__), directory)
        
        if not os.path.exists(base_path):
            os.makedirs(base_path)  # 如果目录不存在则创建
            logger.info(f"创建音乐目录: {base_path}")
        
        for root, _, files in os.walk(base_path):
            for file in files:
                if file.lower().endswith(('.mp3', '.wav')):  # 增加对WAV格式的支持
                    full_path = os.path.join(root, file)
                    file_size = Path(full_path).stat().st_size
                    if file_size == 0:  # 检查文件是否为空
                        logger.warning(f"跳过空文件: {file}")
                        continue
                        
                    music_files.append({
                        "name": file,
                        "path": full_path,
                        "size": file_size,
                        "format": os.path.splitext(file)[1][1:]
                    })
        
        global current_playlist
        current_playlist = music_files
        logger.info(f"扫描到 {len(music_files)} 个音乐文件")
        return {
            "success": True, 
            "files": music_files,
            "directory": base_path,
            "supported_formats": [".mp3", ".wav"]
        }
    except Exception as e:
        logger.error(f"扫描音乐目录失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def play_next() -> dict:
    """播放下一首歌"""
    try:
        global current_track_index
        if not current_playlist:
            return {"success": False, "error": "播放列表为空"}
        
        current_track_index = (current_track_index + 1) % len(current_playlist)
        track = current_playlist[current_track_index]
        return play_music(track["path"])
    except Exception as e:
        logger.error(f"播放下一首失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def play_previous() -> dict:
    """播放上一首歌"""
    try:
        global current_track_index
        if not current_playlist:
            return {"success": False, "error": "播放列表为空"}
        
        current_track_index = (current_track_index - 1) % len(current_playlist)
        track = current_playlist[current_track_index]
        return play_music(track["path"])
    except Exception as e:
        logger.error(f"播放上一首失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def get_current_playlist() -> dict:
    """获取当前播放列表状态"""
    try:
        return {
            "success": True,
            "playlist": current_playlist,
            "current_index": current_track_index,
            "total_tracks": len(current_playlist),
            "current_track": current_playlist[current_track_index] if current_track_index >= 0 and current_playlist else None
        }
    except Exception as e:
        logger.error(f"获取播放列表失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def play_track_by_index(index: int) -> dict:
    """通过索引播放指定歌曲"""
    try:
        if not current_playlist:
            return {"success": False, "error": "播放列表为空"}
        
        if not 0 <= index < len(current_playlist):
            return {"success": False, "error": "无效的索引值"}
            
        global current_track_index
        current_track_index = index
        track = current_playlist[index]
        return play_music(track["path"])
    except Exception as e:
        logger.error(f"播放指定歌曲失败: {str(e)}")
        return {"success": False, "error": str(e)}

# 添加Windows应用程序控制工具
@mcp.tool()
def launch_application(app_path: str) -> dict:
    """启动Windows应用程序。"""
    try:
        subprocess.Popen(app_path)
        logger.info(f"启动应用程序: {app_path}")
        return {"success": True, "message": f"应用程序已启动: {app_path}"}
    except Exception as e:
        logger.error(f"启动应用程序失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def list_running_applications() -> dict:
    """列出所有正在运行的应用程序。"""
    try:
        apps = []
        for proc in psutil.process_iter(['pid', 'name', 'status']):
            try:
                proc_info = proc.info
                apps.append({
                    'pid': proc_info['pid'],
                    'name': proc_info['name'],
                    'status': proc_info['status']
                })
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
        
        logger.info(f"获取到 {len(apps)} 个运行中的应用程序")
        return {"success": True, "applications": apps}
    except Exception as e:
        logger.error(f"获取应用程序列表失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def terminate_application(pid: int) -> dict:
    """终止指定的应用程序。"""
    try:
        process = psutil.Process(pid)
        process.terminate()
        logger.info(f"终止应用程序: PID {pid}")
        return {"success": True, "message": f"应用程序已终止: PID {pid}"}
    except Exception as e:
        logger.error(f"终止应用程序失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def focus_window(window_title: str) -> dict:
    """将指定标题的窗口置于前台。"""
    try:
        def callback(hwnd, windows):
            if win32gui.IsWindowVisible(hwnd):
                title = win32gui.GetWindowText(hwnd)
                if window_title.lower() in title.lower():
                    windows.append(hwnd)
        
        windows = []
        win32gui.EnumWindows(callback, windows)
        
        if not windows:
            return {"success": False, "error": "未找到指定窗口"}
            
        hwnd = windows[0]
        win32gui.ShowWindow(hwnd, win32con.SW_RESTORE)
        win32gui.SetForegroundWindow(hwnd)
        
        logger.info(f"窗口已置于前台: {window_title}")
        return {"success": True, "message": f"窗口已置于前台: {window_title}"}
    except Exception as e:
        logger.error(f"置于前台失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def list_windows() -> dict:
    """列出所有可见窗口。"""
    try:
        def callback(hwnd, windows):
            if win32gui.IsWindowVisible(hwnd):
                title = win32gui.GetWindowText(hwnd)
                if title:
                    windows.append({
                        'handle': hwnd,
                        'title': title
                    })
        
        windows = []
        win32gui.EnumWindows(callback, windows)
        
        logger.info(f"获取到 {len(windows)} 个可见窗口")
        return {"success": True, "windows": windows}
    except Exception as e:
        logger.error(f"获取窗口列表失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def lock_computer() -> dict:
    """锁定Windows计算机屏幕"""
    try:
        user32 = ctypes.windll.User32
        user32.LockWorkStation()
        logger.info("计算机已锁屏")
        return {"success": True, "message": "计算机已锁屏"}
    except Exception as e:
        logger.error(f"锁屏失败: {str(e)}")
        return {"success": False, "error": str(e)}

@mcp.tool()
def unlock_computer(password: str) -> dict:
    """解锁Windows计算机屏幕，需要提供密码"""
    try:
        # 使用Win32API模拟按键
        shell = win32com.client.Dispatch("WScript.Shell")
        # 等待系统准备就绪
        sleep(1)
        # 激活锁屏界面
        shell.SendKeys(" ")
        sleep(1)
        # 输入密码
        shell.SendKeys(password)
        sleep(0.5)
        # 按回车确认
        shell.SendKeys("{ENTER}")
        
        logger.info("已尝试解锁计算机")
        return {"success": True, "message": "已尝试解锁计算机"}
    except Exception as e:
        logger.error(f"解锁失败: {str(e)}")
        return {"success": False, "error": str(e)}



# 启动服务器
if __name__ == "__main__":
    try:
        # 启动MCP服务器
        mcp.run(transport="stdio")
    except KeyboardInterrupt:
        logger.info("程序正在退出...")
    except Exception as e:
        logger.error(f"程序异常退出: {str(e)}")